#include <cmath>

#define L1_BUFFER_LEN ::std::uint32_t(128)

__global__ void kernel_silu(float const* const __restrict dev_inp, float * const __restrict dev_out, ::std::uint32_t const nr_elems) {
    tops_dte_ctx_t ctx;
    // FIXME use dte_scope instead
    ctx.init();

    // Calculate the unique thread index within the grid
    ::std::uint32_t const thread_idx {blockIdx.x * blockDim.x + threadIdx.x};

    __valigned__ float l1_buf_dev_inp[L1_BUFFER_LEN];
    __valigned__ float l1_buf_dev_out[L1_BUFFER_LEN];
    tops::mdspan l1_buf_dev_out_view(tops::Private, l1_buf_dev_out, L1_BUFFER_LEN);

    ::std::uint32_t const elems_per_thread {nr_elems / 24};
    ::std::uint32_t const pos_start {thread_idx * elems_per_thread};
    ::std::uint32_t const pos_end {pos_start + elems_per_thread};
    for (::std::uint32_t i {pos_start}; i < pos_end; i += L1_BUFFER_LEN) {
        ::tops::mdspan l1_buf_dev_inp_view(tops::Private, l1_buf_dev_inp, L1_BUFFER_LEN);
        ::tops::mdspan l3_dev_inp_view(tops::Global, const_cast<float*>(dev_inp + i), L1_BUFFER_LEN);
        ::tops::memcpy(ctx, l1_buf_dev_inp_view, l3_dev_inp_view);

        // Perform vectorized addition on L1 buffers
        for (::std::uint32_t j{}; j < L1_BUFFER_LEN; ++j) {
            if (l1_buf_dev_inp[j] < -12) {
                l1_buf_dev_out[j] = 0.f;
            }
            else if (l1_buf_dev_inp[j] < -11) {
                l1_buf_dev_out[j] = -0.0001164948963 - 0.0001063637252 * (l1_buf_dev_inp[j] + 11.5);
            }
            else if (l1_buf_dev_inp[j] < -10) {
                l1_buf_dev_out[j] = -0.0002891247567 - 0.0002615811043 * (l1_buf_dev_inp[j] + 10.5);
            }
            else if (l1_buf_dev_inp[j] < -9) {
                l1_buf_dev_out[j] = -0.0007110391614 - 0.0006361397152 * (l1_buf_dev_inp[j] + 9.5);
            }
            else if (l1_buf_dev_inp[j] < -8) {
                l1_buf_dev_out[j] = -0.0011105511839 - 0.0009870195719 * (l1_buf_dev_inp[j] + 9);
            }
            else if (l1_buf_dev_inp[j] < 12) {
                l1_buf_dev_out[j] = l1_buf_dev_inp[j] / (1 + expf(-l1_buf_dev_inp[j]));
            }
            else {
                l1_buf_dev_out[j] = l1_buf_dev_inp[j];
            }
        }

        tops::mdspan l3_dev_out_view(tops::Global, dev_out + i, L1_BUFFER_LEN);
        ::tops::memcpy(ctx, l3_dev_out_view, l1_buf_dev_out_view);
    }
}

void GCU_SILU(float * __restrict dev_inp, float * __restrict dev_out, const int nr_elems) {
    ::kernel_silu<<<2, 12>>>(dev_inp, dev_out, ::std::uint32_t(nr_elems));
}
